import os
import sys
import cPickle
import time
import subprocess
from math import log10, log

class memoized(object):
   """Decorator that caches a function's return value each time it is called.
   If called later with the same arguments, the cached value is returned, and
   not re-evaluated.
   """
   def __init__(self, func):
      self.func = func
      self.cache = {}
   def __call__(self, *args):
      try:
         return self.cache[args]
      except KeyError:
         value = self.func(*args)
         self.cache[args] = value
         return value
      except TypeError:
         # uncachable -- for instance, passing a list as an argument.
         # Better to not cache than to blow up entirely.
         return self.func(*args)
   def __repr__(self):
      """Return the function's docstring."""
      return self.func.__doc__
   def __get__(self, obj, objtype):
      """Support instance methods."""
      return functools.partial(self.__call__, obj)


def path_dict(genome, fpath):
    pathDB= cPickle.load(open(os.path.join(fpath, 'pathDB.pkl'), 'rb'))
        
    return pathDB[genome]


basecomp= {'a':'t', 'c':'g', 't':'a', 'g':'c', 'n':'n',\
               'A':'T', 'C':'G', 'T':'A', 'G':'C', 'N':'N'}

def complement(s):
    comps= [basecomp[base] for base in s]
    if type(s) is list:
        return comps
    return ''.join(comps)


def revcomp(s):
    """
    if type(s) is list:
        s.reverse()
        return complement(s)
    # if it is a string
    """
    s= s[::-1]
    return complement(s)


def GC(s):
    s= s.upper()
    gc= (s.count('G')+s.count('C'))*100.0/len(s)
    return gc


# all values are negative (Allawani H_Biochemistry_1997)
deltaHSG= {'AA': [-7.9, -21.3, -1.0],\
               'TT': [-7.9, -21.3, -1.0],\
               'AT': [-7.2, -20.4, -0.88],\
               'TA': [-7.2, -21.3, -0.58], \
               'CA': [-8.5, -22.7, -1.45], \
               'TG': [-8.5, -22.7, -1.45], \
               'GT': [-8.4, -22.4, -1.44], \
               'AC': [-8.4, -22.4, -1.44], \
               'CT': [-7.8, -21.0, -1.28], \
               'AG': [-7.8, -21.0, -1.28], \
               'GA': [-8.2, -22.2, -1.30], \
               'CG': [-10.6, -27.2, -2.17], \
               'GC': [-9.8, -24.4, -2.24], \
               'GG': [-8.0, -19.9, -1.84], \
               'CC': [-8.0, -19.9, -1.84], \
               'TC': [-8.2, -22.2, -1.30]}

deltaG_MM= {'G':{'AA':0.17,'AC':0.81,'AG':-0.25,'AT':-1.30,'CA':0.47,'CC':0.79,\
                     'CG':-2.24,'CT':0.62,'GA':-0.52,'GC':-1.84,'GG':-1.11,'GT':0.08,\
                     'TA':-1.44,'TC':0.98,'TG':-0.59,'TT':0.45},\
            'C':{'AA':0.43,'AC':0.75,'AG':0.03,'AT':-1.45,'CA':0.79,'CC':0.70,\
                     'CG':-1.84,'CT':0.62,'GA':0.11,'GC':-2.17,'GG':-0.11,'GT':-0.47,\
                     'TA':-1.28,'TC':0.40,'TG':-0.32,'TT':-0.12},\
            'A':{'AA':0.61,'AC':0.88,'AG':0.14,'AT':-1.0,'CA':0.77,'CC':1.33,\
                     'CG':-1.44,'CT':0.64,'GA':0.02,'GC':-1.28,'GG':-0.13,'GT':0.71,\
                     'TA':-0.88,'TC':0.73,'TG':0.07,'TT':0.69},\
            'T':{'AA':0.6,'AC':0.92,'AG':0.42,'AT':-0.58,'CA':1.33,'CC':1.05,'CG':-1.30,\
                     'CT':0.97,'GA':0.74,'GC':-1.45,'GG':0.44,'GT':0.43,\
                     'TA':-1.0,'TC':0.75,'TG':0.34,'TT':0.68}}

def dimer(F, R, pntcutoff=2.5):

    # this means a 50% difference in deltaG from a full-match
    if pntcutoff>5: 
        pntcutoff= 2.5

    F= F.upper() # which is given as 5`-3`
    R= R.upper() # which is given as 5`-3`
    FdeltaG= []
    RdeltaG= []
    if F == R:
        # is the terminal A-T, G-C penalty necessary for dimer purposes
        # it is difficult to calculate it for dimers 
        #Fterm= F[0]+F[-1]
        #FdeltaG= 1.03*(Fterm.count('A')+Fterm.count('T'))+0.98*(Fterm.count('G')+Fterm.count('C'))
        [FdeltaG.append(olcount(F, keys) * deltaHSG[keys][2]) for keys in deltaHSG]
        # calculate mean deltaG
        MdeltaG= sum(FdeltaG)
        #print MdeltaG

    elif F != R:
        # is the terminal A-T, G-C penalty necessary for dimer purposes
        #Fterm= F[0]+F[-1]
        #FdeltaG= 1.03*(Fterm.count('A')+Fterm.count('T'))+0.98*(Fterm.count('G')+Fterm.count('C'))
        #Rterm= R[0]+R[-1]
        #RdeltaG= 1.03*(Rterm.count('A')+Rterm.count('T'))+0.98*(Rterm.count('G')+Rterm.count('C'))

        [(FdeltaG.append(olcount(F, keys) * deltaHSG[keys][2]), RdeltaG.append(olcount(R, keys) * deltaHSG[keys][2])) for keys in deltaHSG]
        
        # calculate mean deltaG
        MdeltaG= (sum(FdeltaG) + sum(RdeltaG)) / 2
        #print MdeltaG

    # F: NN-5`---------------------3`
    #                    | || || ||
    # R:           << 3`------------------------5`-NN
    # For the first round add N to 5` of either F or R
    if len(F) - len(R) > 0:
        longest= len(F)
        R= 'N'*(len(F)-len(R))+R
    elif len(F) - len(R) < 0:
        longest= len(R)
        F= 'N'*(len(R)-len(F))+F
    elif len(F) == len(R):
        longest= len(F)
    R= R[::-1] # reverse so that it is 3`-5` (do a string reverse)
    J= 6
    while J < longest: 
        listF= F[-J:]
        listR= R[:J]
        deltaG= 0
        I= 0
        while I < len(listF)-1:
            if listF[I] == complement(listR[I]) and listF[I] != "N" and listF[I+1] != "N" and listR[I+1] != "N":
                deltaG+= deltaG_MM[listF[I]][listF[I+1]+listR[I+1]]
            I+=1
        if deltaG < MdeltaG*pntcutoff*0.20: # if delta G is larger than
            #print '1R: dimer deltaG is %s with 5`-%s and 3`-%s' %(deltaG, listF, listR)
            return 'FAIL'
        J+=1

    if F.find('N') > -1: 
        F.replace('N','')
    elif R.find('N') > -1: 
        R.replace('N','')

    # F:                 5`---------------------3`-NN
    #                       | || || ||
    # R:  NN-3`------------------------5`>>
    # For the second round add N to 3` of either F or R
    # don't forget that R is already reversed
    if len(F) - len(R) > 0:
        longest= len(F)
        R= 'N'*(len(F)-len(R))+R # add to the begining (R is in reverse so the 3`)
    elif len(F) - len(R) < 0:
        longest= len(R)
        F= F+'N'*(len(R)-len(F)) # add to the end the 3`
    elif len(F) == len(R):
        longest= len(F)
    J= 6
    while J < longest: 
        listF= F[:J]
        listR= R[-J:]
        deltaG= 0
        I= 0
        while I < len(listF)-1:
            if listF[I] == complement(listR[I]) and listF[I] != "N" and listF[I+1] != "N" and listR[I+1] != "N":
                deltaG+= deltaG_MM[listF[I]][listF[I+1]+listR[I+1]]
            I+=1
        if deltaG < MdeltaG*pntcutoff*0.20: # if delta G is larger than
            #print '2R: dimer deltaG is %s with 5`-%s and 3`-%s' %(deltaG, listF, listR)
            return 'FAIL'
        J+=1
    return "PASS"



#advanced Tm calculation
#nearest-neighbour thermedynamics calculation of melting temperature
#corrected for monovalent ion, Mg+2, dNTP and even DMSO concentrations
#Credits:  
#Main author: Ogan ABAAN 
#Overcount function: Greg Singer <singerg@tcd.ie>


def olcount(s,pattern):
    """Returns how many pattern on s, works for overlapping"""
    count= 0
    x= 0
    while 1:
        try:
            i= s.index(pattern,x)
        except ValueError:
            break
        count+= 1
        x= i+1
    return count

# deltaH (2nd column unit in kcal/mol) and deltaS (3rd column unit in cal/mol K) NN parameters in 1M NaCl (SantaLucia J, PNAS, 1998)
deltaHS= {'AA': [-7.9, -22.2],\
              'TT': [-7.9, -22.2],\
              'AT': [-7.2, -20.4],\
              'TA': [-7.2, -21.3], \
              'CA': [-8.5, -22.7], \
              'TG': [-8.5, -22.7], \
              'GT': [-8.4, -22.4], \
              'AC': [-8.4, -22.4], \
              'CT': [-7.8, -21.0], \
              'AG': [-7.8, -21.0], \
              'GA': [-8.2, -22.2], \
              'CG': [-10.6, -27.2], \
              'GC': [-9.8, -24.4], \
              'GG': [-8.0, -19.9], \
              'CC': [-8.0, -19.9], \
              'TC': [-8.2, -22.2]}


def Tm(s, dna, Na, K, Tris, Mg, dNTP, DMSO):
    
    #dna= TMdict[dna]
    #Na= TMdict[Na]
    #K= TMdict[K]
    #Tris= TMdict[Tris]
    #Mg= TMdict[Mg]
    #dNTP= TMdict[dNTP]
    #DMSO= TMdict[DMSO]

    # dna      the oligo concentration (nM)
    # Na       the Na+ concentration (mM) 
    # K        the K+ concetration (mM)
    # Tris     the Tris+ concentration (mM)
    # Mg       the Mg+2 concentration (mM)
    # dNTP     the dNTP concentration (mM)
    # DMSO     correction for DMSO (%)

    s= s.upper()
    # Molar concentration of all monovalent ions (Owczarzy R, Biochemistry, 2008)
    Mon= (Na+K+Tris/2.0)/1000.0

    # Molar concentration for non-self complementary oligonucleotide (SantaLucia J, PNAS, 1998)
    Ct= (dna/4.0)*1e-9

    # Molar concentration of free Mg+2 (Owczarzy R, Biochemistry, 2008)    
    fMg= (Mg-dNTP)/1000.0

    if Mon > 0:
        ionR= fMg**0.5/Mon
        #print 'Ion ratio is', ionR

    # calculate faction of GC    
    fGC= GC(s)/100.0    

    # DMSO correction (von Ahsen N, Clin Chem, 2000)
    # take the average of all four coefficients previously published
    DMSO_coeff= round((0.75 + 0.6 + 0.675 +0.5)/4,2)
    # 1% DMSO reduces Tm by DMSO_coeff
    dmso= DMSO_coeff*DMSO

    # double check for the 'fit-model' for range of criteria
    if Mon == 0 or Mg == 0:
        pass
    elif 1e-10 <= (dna*1e-9) <= 0.1 and 7 < len(s) < 61 and 0.015 <= Mon <= 1.2 and 0.01 <= Mg < 600 and dNTP <= (Mg*1.2):
        pass
    else:
        print 'buffer/DNA conditions out of range for calculation'
        sys.exit()


    dh, ds = [], []
    # terminal base corrections
    if s[0]=='G' or s[0]=='C':
        dh.append(0.1)
        ds.append(-2.8)
    elif s[0]=='A' or s[0]=='T':
        dh.append(2.3)
        ds.append(4.1)
    if s[-1]=='G' or s[-1]=='C':
        dh.append(0.1)
        ds.append(-2.8)
    elif s[-1]=='A' or s[-1]=='T':
        dh.append(2.3)
        ds.append(4.1)

    [(dh.append(olcount(s, keys) * deltaHS[keys][0]), ds.append(olcount(s, keys) * deltaHS[keys][1])) for keys in deltaHS]
    
    # print dh, ds
    Rconst= 1.987 #gas constant in cal/mol K 
    tmNN= (1000 * sum(dh))/(sum(ds) + (Rconst*log(Ct))) #result is in Kelvin
    # print tmNN-273.15

    # Decision tree based on Owczarzy R, Biochemistry, 2008
    if Mon == 0 or (Mon != 0 and ionR >= 6.0):
    # incorporate Mg correction (Owczarzy R, Biochemistry, 2008)    
        #various constants (1/K)
        Va= 3.92e-5
        Vb= -9.11e-6
        Vc= 6.26e-5
        Vd= 1.42e-5
        Ve= -4.82e-4
        Vf= 5.25e-4
        Vg= 8.31e-5

        tmMg= (1.0/tmNN)+Va+(Vb*log(fMg))+(fGC*(Vc+Vd*log(fMg)))+((Ve+Vf*log(fMg)+(Vg*(log(fMg))**2))/(2*(len(s)-1)))
        tmMg= (1.0/tmMg)-273.15
        return tmMg-dmso

    elif Mon != 0 and ionR < 0.22:
    # incorporate salt correction (for all monovalent ions) (OwczarzyR, Biochemistry, 2004)
        tmMon= (1.0/tmNN)+(4.29*fGC-3.95)*1e-5*log(Mon)+9.40*1e-6*log(Mon)**2
        tmMon= (1.0/tmMon)-273.15
        return tmMon-dmso

    elif Mon != 0 and 0.22 <= ionR < 6.0:
    # incorporate Mg correction (OwczarzyR, Biochemistry, 2008)    
        #various constants (1/K) corrected for Mon+ as well 
        Va= 3.92e-5*(0.843-(0.352*Mon**0.5*log(Mon)))
        Vb= -9.11e-6
        Vc= 6.26e-5
        Vd= 1.42e-5*(1.279-4.03e-3*log(Mon)-8.03e-3*log(Mon)**2)
        Ve= -4.82e-4
        Vf= 5.25e-4
        Vg= 8.31e-5*(0.486-0.258*log(Mon)+5.25e-3*log(Mon)**3)

        tmMg= (1.0/tmNN)+Va+(Vb*log(fMg))+(fGC*(Vc+Vd*log(fMg)))+((Ve+Vf*log(fMg)+(Vg*(log(fMg))**2))/(2*(len(s)-1)))
        tmMg= (1.0/tmMg)-273.15
        return tmMg-dmso

    else:
        print "error in Tm calculation"

def gethms(seconds):
    seconds= int(seconds)
    days= seconds/(3600*24)
    seconds-= 24*3600*days
    hours= seconds/3600
    seconds-= 3600*hours
    minutes= seconds/60
    seconds-= 60*minutes
    if days==0:
        if hours==0:
            if minutes==0: 
                return "%ds" % (seconds)
            else: 
                return "%dm:%ds" % (minutes, seconds)
        else:
            return "%dh:%dm:%ds" % (hours, minutes, seconds)
    else:
        return "%dd:%dh:%dm:%ds" % (days, hours, minutes, seconds)

# establish some window properties, only once
if sys.platform == 'darwin':
    try:
        row, column = os.popen('stty size','r').read().split()
        column= int(column)
    except OSError:
        column= 80
elif sys.platform == 'linux2':
    column= 80
### end

def progress_info(i, N, elapsed, decimal= False, silence= 1):
    if not silence:
        return
    
    if i == 0:
        statement= '0%% completed in %s - ETA: unknown' % (gethms(elapsed))
    # i= completed
    # N= total
    # elapsed= elapsed time
    # mode= 1 or 2 for no decimal and decimal place in output
    else:
        if decimal==True:
            statement= '%.1f%% completed in %s - ETA: %s' % ((i*100.0/N), gethms(elapsed), gethms((1.0*N*elapsed/i)-elapsed))
        else:
            statement= '%d%% completed in %s - ETA: %s' % ((i*100/N), gethms(elapsed), gethms((1.0*N*elapsed/i)-elapsed))

    sys.stdout.write('\r%s' % ''.join([statement,' '*(column-len(statement))]))
    sys.stdout.flush()
    return

def synopsis(text, silence= 1):
    if not silence:
        return
    sys.stdout.write('\r%s' % ''.join([text, ' '*(column-len(text))]))
    sys.stdout.flush()
    return


def DBgenome(gpath, CHR, start, stop= 0):
    input= open(os.path.join(gpath, '%s.fa' % CHR))
    header= len(input.readline()) # read the header line

    input.seek(header + (start - 1) / 50 + (start - 1))
    if stop == 0:
        # read only one byte
        return input.read(1).upper()
    else:
        lenS= stop - start + 1
        seq = "".join([input.readline().strip() for i in xrange(lenS/50 + 2)])
        res = seq[:lenS].lower()
        return res

def mean(s):
    total= sum([float(p) for p in s])
    return total/len(s)

def stdev(s):
    diff= 0
    avg= mean(s)
    for items in s:
        diff+= (float(items)-avg)**2
    return (diff/len(s))**0.5


def metric_system(s):
    conv= {18:['exa', 'E'], 15:['peta', 'P'], 12:['tera', 'T'],\
           9:['giga', 'G'], 6:['mega', 'M'], 3:['kilo', 'K'],\
           -3:['milli', 'm'], -6:['micor', 'u'], -9:['nano', 'n'],\
           -12:['pico', 'p'], -15:['femto', 'f'], -18:['atto', 'a'],\
           0:['', '']}

    # find the closest value
    value= (int(log10(s)) / 3) * 3

    return "%.3f %s" %(s/(10.0**value), conv[value][1])
